import torch
from torch import nn, Tensor
from typing import Iterable, Dict


class CosineSimilarityLossWithTemperature(nn.Module):
    """
    Dot product or cosine similarity
    Reference: https://github.com/princeton-nlp/SimCSE/blob/511c99d4679439c582beb86a0372c04865610b6b/simcse/models.py#L35

    """

    def __init__(self, temp = 0.02, loss_fct = nn.MSELoss()):
        super().__init__()
        self.temp = temp
        self.cos = nn.CosineSimilarity(dim=-1)

    def forward(self, sentence_features: Iterable[Tensor], labels: Tensor):
        output = self.cos(sentence_features[0], sentence_features[1]) / self.temp
        return self.loss_fct(output, labels.view(-1))

